package com.olive.filter;

import lombok.extern.slf4j.Slf4j;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;

@Slf4j
public class GzipResponseWrapper extends HttpServletResponseWrapper {

    /**
     * 字节数组缓冲流，用来保存截获到的输出数据
     */
    private ByteArrayOutputStream buffer;

    /**
     * 重新定义servlet输出流，改变输出目的地将响应内容输出到给定的字节数组缓冲流中
     */
    private GzipResponseWrapper.CustomServletOutputStream servletOutputStream;

    /**
     * 同上
     */
    private PrintWriter writer;

    public GzipResponseWrapper(HttpServletResponse response) {
        super(response);
        //original HttpServletResponse object
        buffer = new ByteArrayOutputStream();
        servletOutputStream = new GzipResponseWrapper.CustomServletOutputStream(buffer);
        try {
            writer = new PrintWriter(new OutputStreamWriter(buffer, response.getCharacterEncoding()), true);
        } catch (UnsupportedEncodingException e) {
            log.error("GZipHttpServletResponse", e);
        }
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return servletOutputStream;
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        return writer;
    }

    @Override
    public void flushBuffer() throws IOException {
        if (servletOutputStream != null) {
            servletOutputStream.flush();
        }
        if (writer != null) {
            writer.flush();
        }
    }

    /**
     * 向外部提供一个获取截获数据的方法
     * @return 从response输出流中截获的响应数据
     */
    public byte[] getOutputData() throws IOException {
        flushBuffer();
        return buffer.toByteArray();
    }

    private static class CustomServletOutputStream extends ServletOutputStream {

        /**
         * 字节数组缓冲流，用来保存截获到的输出数据
         */
        private ByteArrayOutputStream buffer;

        public CustomServletOutputStream(ByteArrayOutputStream buffer) {
            this.buffer = buffer;
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setWriteListener(WriteListener listener) {
        }

        /**
         * 重写输出流相关的方法
         * 将输出数据写出到给定的ByteArrayOutputStream缓冲流中保存起来
         * @param b 输出的数据
         * @throws IOException
         */
        @Override
        public void write(int b) throws IOException {
            buffer.write(b);
        }
    }
}